import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.model_selection import train_test_split
from sklearn.svm import LinearSVC
from sklearn.metrics import classification_report

# 假设你已经有一个包含文本数据和对应标签的数据集
# 这里只是一个简单的示例，实际数据集可能更复杂
data = {
    'text': [
        "SELECT * FROM users WHERE username = 'admin' AND password = '123';",  # SQL注入
        "<script>alert('XSS attack!');</script>",  # XSS攻击
        "rm -rf /",  # Shell注入
        "GET /index.html HTTP/1.1",  # 正常流量
        # 可以根据需要添加更多样本
    ],
    'label': ['SQL', 'XSS', 'Shell', 'Normal']
}

df = pd.DataFrame(data)

# 特征提取
vectorizer = TfidfVectorizer()
X = vectorizer.fit_transform(df['text'])

# 将标签转换为数值
label_dict = {'SQL': 0, 'XSS': 1, 'Shell': 2, 'Normal': 3}
y = df['label'].map(label_dict)

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 训练模型
model = LinearSVC()
model.fit(X_train, y_train)

# 预测
y_pred = model.predict(X_test)

# 评估模型性能
# 评估模型性能
print(classification_report(y_test, y_pred, zero_division=1, labels=['SQL', 'XSS', 'Shell', 'Normal']))

